import argparse
import torch
from mkd_model import MKD_Trainer, ResNet18ForTabular
from attack_evaluator import MembershipInferenceAttacker
from torchvision.models import resnet18


def parse_args():
    parser = argparse.ArgumentParser(description="MKD评估入口")
    parser.add_argument("--dataset", type=str, default="cifar100", choices=["cifar100", "texas100", "purchase100"],
                        help="数据集名称")
    parser.add_argument("--model", type=str, default="resnet18", choices=["resnet18", "alexnet"], help="模型类型")
    parser.add_argument("--ckpt_path", type=str, default="E:/xt/雷达/培训/code/checkpoints/teacher_cifar100_resnet18_final.pth",
                        help="目标模型权重路径")
    return parser.parse_args()


def load_target_model(dataset_name, model_type, ckpt_path, device):
    """加载目标模型（MKD学生模型）"""
    if dataset_name == "cifar100":
        model = resnet18(pretrained=False, num_classes=100).to(device)
    else:
        # 表格数据：先获取特征维度
        from data_loader import get_data_loaders
        dt_loader, _, _ = get_data_loaders(dataset_name, batch_size=512)
        input_dim = next(iter(dt_loader))[0].shape[1]
        model = ResNet18ForTabular(input_dim, num_classes=100).to(device)

    # 加载权重
    model.load_state_dict(torch.load(ckpt_path, map_location=device))
    model.eval()
    return model


if __name__ == "__main__":
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 1. 加载数据
    from data_loader import get_data_loaders

    dt_loader, ds_loader, test_loader = get_data_loaders(args.dataset, batch_size=512)

    # 2. 加载目标模型
    target_model = load_target_model(args.dataset, args.model, args.ckpt_path, device)

    # 3. 初始化攻击者并运行攻击评估
    attacker = MembershipInferenceAttacker(
        target_model=target_model,
        dt_loader=dt_loader,  # 成员数据（D_T）
        test_loader=test_loader,  # 非成员数据（测试集）
        device=device
    )
    attack_results = attacker.run_all_attacks()

    # 4. 评估模型任务准确率（最终报告用）
    from mkd_model import MKD_Trainer

    mkd_trainer = MKD_Trainer(dataset_name=args.dataset, model_type=args.model)
    test_acc = mkd_trainer.evaluate_model(target_model, test_loader)
    print(f"目标模型（MKD学生模型）测试准确率：{test_acc:.4f}")